module ar::FlowGraphsAndClassDiagrams

import lang::ofg::ast::FlowLanguage;
import lang::ofg::ast::Java2OFG;
import List;
import Relation;
import lang::java::m3::Core;

import IO;
import vis::Figure; 
import vis::Render;

public loc emptyId = |id:///|;
alias OFG = rel[loc from, loc to];

OFG buildGraph(Program p)
	// constructor is: loc id, list[loc] formalParameters.
	// newAssign is: loc target, loc class, loc constructor, list[loc] actualParameters.
	// index() returns a list of all legal index values for a given list.

	// add flows with the form "s = new Student("Smith");" where actual parameter "Smith" flows to formal parameter "studentName".
  = { <actualParams[i], formalParams[i]> | newAssign(target, class, ctor, actualParams) <- p.statements, constructor(ctor, formalParams) <- p.decls, i <- index(actualParams) }
	// add flows with the form "findBook("The Selfish Gene");" where actual parameter "The Selfish Gene" flows to formal parameter "bookName".
  + { <actualParams[i], formalParams[i]> | call(target, cast, receiver, func, actualParams) <- p.statements, method(func, formalParams) <- p.decls, i <- index(actualParams) }
  	// add flows with the form "s = new Student("Smith");" where "Student.Student.this" flows to "s".
  + { <ctor + "this", target> | newAssign(target, class, ctor, actualParams) <- p.statements }
  	// add flows with the form "b = c.findBook("The Selfish Gene");" where "c" flows to "findBook.this".
  + { <receiver, func + "this"> | call(target, cast, receiver, func, actualParams) <- p.statements }
  	// add flows with the form "b = c.findBook("The Selfish Gene");" where "Book.findBook.return" flows to "b".
  + { <func + "return", target> | call(target, cast, receiver, func, actualParams) <- p.statements, target != emptyId }
  	// add flows with the form "a = b", where "b" flows to "a".
  + { <source, target> | assign(target, cast, source) <- p.statements }
  ;
  
rel[loc super, loc sub] getGeneralizationRelations(Program p, M3 myModel)
	// look for the keyword "extends".
  = { <superclass, subclass> | <subclass, superclass> <- myModel@extends, superclass in classes(myModel), subclass in classes(myModel) }
;

rel[loc,loc] getRealizationRelations(Program p, M3 myModel)
	// look for the keyword "implements".
  = { <superclass, subclass> | <subclass, superclass> <- myModel@implements, superclass in classes(myModel), subclass in classes(myModel)  }
;
  
// Associations are already covered by the OFG, so we do not need to write a specific method for that relation.

rel[loc,loc] getNestedClassRelations(Program p, M3 myModel)
	// Look for cases where class B is a nested class inside class A. Nested classes definitely depend on their containing class.
  = { <c, n> | c <- classes(myModel), n <- nestedClasses(myModel, c) }
;

rel[loc,loc] buildGenSet(Program p, bool forward)
	// Only nodes of type cs.this have a non-empty gen set, so we add gens of the form "Student.Student.this -> Student"
	
  	// add flows with the form "s = new Student("Smith");" where "Student.Student.this" generates type "Student".
  = forward ? ({ <ctor + "this", class> | newAssign(target, class, ctor, actualParams) <- p.statements })
  			: 	(	// add flows with the form "x = (C)y", store gen[y] = {C}
  					{ <source, cast> | assign(target, cast, source) <- p.statements, cast != emptyId }
  					// add flows with the form "x = (C)y.method(a1,...,ak);", store gen[m.return] = {C}
  				  + { <func + "return", cast> | call(target, cast, receiver, func, actualParams) <- p.statements, cast != emptyId  }
  				  // Note that in these cases () denotes the type coercion operator, also known as "casting".
  				)
;

OFG prop(OFG graph, rel[loc,loc] GEN, rel[loc,loc] KILL, bool forward) {
	// This statement switches the order of the graph tuples from (from,to) to (to,from)
	// Note that it may also change the "order" of the set, since sets don't have an order.
	graph_reverse = graph<to,from>;
	// Here the inline functions pred and succ are defined.
	// Predecessor(n) will return all the <from> parts of the relation tuples where <to>=n
	set[loc] pred(loc n) = graph_reverse[n];
	// Successor(n) will return all the <to> parts of the relation tuples where <from>=n
	set[loc] succ(loc n) = graph[n];
	// First initialize IN as the empty set.
	OFG IN = {};
	// Now generate OUT based on parameters gen and kill (IN is still empty of course).
	OFG OUT = GEN + (IN - KILL);
  	// Solve is a statement that repeats a certain computation as long as it causes changes. So basically this is the loop of the propagation algorithm.
	solve (IN, OUT) {
  		// Possible bug: when back is FALSE (so we do NOT have back propagation, but p <- succ(n), \o <- OUT[p] IS back propagation!
		// This means that forward propagation is used when back is TRUE. Seems wrong! 
		// Note: I fixed above bug by renaming back to forward.
		
		// Carrier takes in a relation set containing tuples and basically turns it into an array. So carrier({<1,10>, <2,20>}) will return {10,20,2,1}
		// The letter 'o' is a keyword in Rascal to denote map/relation composition. The \ denotes the logical not statement, but in this case it is used
		// as the escape character to turn 'o' from a keyword into a plain variable.
		// So as a whole this line loops over all loc's contained anywhere in g, takes the successor/predecessor, takes the OUT set of that successor/predecessor
		// and generates tuples containing a loc from g on the left and a loc of the OUT of the successor/predecessor on the right. Looking up a certain
		// n in IN will then give in IN set of that n.
		// Note: if OUT[p] is empty, then NO tuple for that n is formed.
    	IN = { <n,\o> | n <- carrier(graph), p <- (forward ? pred(n) : succ(n)), \o <- OUT[p] };
    	// Again, calculate the out set: generated data plus incoming data, minus "killed" data.
    	OUT = GEN + (IN - KILL);
  	}
  	return OUT;
}

// Get the class of a location, if the location is a method(or constructor), field, variable or already a class.
set[loc] getClassOfLocation(M3 myModel, loc myLocation)
{
	// if location is already a class, return it.	
	if (isClass(myLocation))
	{ return { myLocation }; }
	
	// location is a variable in a method, return the class that the method is in.
	if (isVariable(myLocation))
	{ return { myClass | myClass <- classes(myModel), myLocation in variables(myModel, myClass) }; }

	// location is a parameter of a method, return the class that the method is in.
  	if (isParameter(myLocation))
  	{ return { myClass | myClass <- classes(myModel), myLocation in parameters(myModel, myClass) }; }

  	// if location is a method, return the class.
  	if (isMethod(myLocation))
  	{ return { myClass | myClass <- classes(myModel), myLocation in methods(myModel, myClass) }; }

  	// if location is a field, return the class.
  	if (isField(myLocation))
  	{ return { myClass | myClass <- classes(myModel), myLocation in fields(myModel, myClass) }; }
	
 	// return in any case.
	return {};
}

// There is probably a better way, but built-in size/isEmpty functions do not accept set[loc].
public bool nonEmpty(set[loc] someSet)
{
	int counter = 0;
	for (s <- someSet)
	{
		counter = counter + 1;
		if (counter > 0)
		{ return true; }
	}
	return false;
}

// Helper method for getClassOfLocation that returns all variables inside the methods(including constructor) of a class.
 public set[loc] variables(M3 m, loc class)
 = { var | method <- elements(m, class), isMethod(method), var <- elements(m, method), isVariable(var) }
 ;
 
 // Helper method for getClassOfLocation that returns all parameters inside the methods(including constructor) of a class.
 public set[loc] parameters(M3 m, loc class)
 = { param | method <- elements(m, class), isMethod(method), param <- elements(m, method), isParameter(param) }
 ;

// Used to simplify the graph if there are many relations/edges going to both superclasses and subclasses.
private rel[loc,loc] LowestCommonAncestor(rel[loc,loc] graph, rel[loc,loc] generalizations, M3 myModel)
{
	allClasses = classes(myModel);
	rel[loc,loc] lcaGraph = graph;
	
	set[loc] getGraphEdges(loc n) = lcaGraph[n];
	set[loc] getGeneralizations(loc n) = generalizations[n];
	
	solve (lcaGraph)
	{
		for (cl1 <- allClasses)
		{
			// For all cases where OUT[some_loc] is for example { User, InternalUser }, we want to eliminate the subclass InternalUser
  			// such that the new OUT[some_loc] is just {User}.
			edges = getGraphEdges(cl1);
			lcaGraph = lcaGraph - { <cl1, g> | e <- edges, g <- getGeneralizations(e), cl1 != g };
			// Secondly, for all cases where OUT[some_loc] is for example { Journal, Book, Report } we want to replace the subclasses
			// with the superclass, such that the new OUT[some_loc] is just { Document }. We *only* want to do this when *all* subclasses
			// of Document are included in the original OUT[some_loc}, so subclasses(Document) = { Journal, Book, Report ). If something misses we do not replace.
			for (cl2 <- allClasses)
			{
				if (cl1 != cl2)
				{
					edges = getGraphEdges(cl1);
					gens = getGeneralizations(cl2);
					if (nonEmpty(edges) && nonEmpty(gens))
					{
						// check if subclasses is a subset of edges
						if (gens <= edges)
						{	// delete the subclasses and replace with the superclass.
							lcaGraph = (lcaGraph - { <cl1, g> | g <- gens }) + { <cl1, cl2> };
						}
					}
				}
			}
		}
	}
	return lcaGraph;
}

public void drawDiagramOfGraph(OFG flowgraph, Program p, M3 m) {
	// find attributes to draw associations.
  	rel[loc, loc] associaton_edges = { <from_class, to> | <from,to> <- flowgraph, from.scheme == "java+field", from_class <- getClassOfLocation(m, from), from_class != to, to in classes(m) };
  	// find local variables and method parameters to draw dependencies.
  	rel[loc, loc] dependency_edges = { <from_class, to> | <from,to> <- flowgraph, (from.scheme == "java+variable" || from.scheme == "java+parameter"), from_class <- getClassOfLocation(m, from), from_class != to, to in classes(m)  };
  	
  	// associations are stronger relations than dependencies, so they take precedence if they overlap.
  	dependency_edges = dependency_edges - associaton_edges;
  	
  	rel[loc,loc] generalizations = getGeneralizationRelations(p, m);
  	rel[loc,loc] realizations = getRealizationRelations(p, m);
  	rel[loc,loc] nestedclasses = getNestedClassRelations(p, m);
  	
  	// Debug output:
  	println("###### DEPENDENCY EDGES ####################");
	for (g <- dependency_edges)
	{
		println(g);
	}
  	println("###### ASSOCIATION EDGES ####################");
	for (g <- associaton_edges)
	{
		println(g);
	}
  	println("###### GENERALIZATION EDGES ####################");
	for (g <- generalizations)
	{
		println(g);
	}

  	//rel[loc,loc] dependencies = getDependencyRelations(p, m);
  	known_relations = generalizations + realizations + nestedclasses;
	
	// If there are associations towards all subclasses of a class, replace those associations with just the class association.
	simplified_associaton_edges = LowestCommonAncestor(associaton_edges, generalizations, m);
	// If there are dependencies towards all subclasses of a class, replace those dependencies with just the class dependency.
	simplified_dependency_edges = LowestCommonAncestor(dependency_edges, generalizations, m);
	
	// removed relations that we already know from the M3, such as generalization, realization, subclasses (aggregation).
	simplified_associaton_edges = simplified_associaton_edges - known_relations;
	simplified_dependency_edges = simplified_dependency_edges - known_relations;
	
	// Take the union of associations and dependencies.
	total_edges = simplified_associaton_edges + simplified_dependency_edges;
	// If a combination of association and dependency relations towards all subclasses of a class, replace those relations with just the class relation.
	simplified_total_edges = LowestCommonAncestor(total_edges, generalizations, m);
	// find all relations that were removed due to the simplification.
	removed = total_edges - simplified_total_edges;
	// find all relations that were added due to the simplification.
	added = simplified_total_edges - total_edges;
	
	// remove all the removed relations from the associations and dependencies.
	simplified_associaton_edges = simplified_associaton_edges - removed;
	simplified_dependency_edges = simplified_dependency_edges - removed;
	
	// all the relations that were added are added to the associations, since they resulted of a combination of dependencies and associations,
	// and we consider associations to be stronger.
	simplified_associaton_edges = simplified_associaton_edges + added;
	
	// remove all relations that come from Main, since the book by Tonella and Potrich ignores them as well.
	// only use this when verifying the code on the example eLib project!
	//simplified_associaton_edges = { <from,to> | <from,to> <- simplified_associaton_edges, from != |java+class:///Main| };
	//simplified_dependency_edges = { <from,to> | <from,to> <- simplified_dependency_edges, from != |java+class:///Main| };
	
	// More debug output:
	println("############ SIMPLIFIED ASSOCIATIONS: ############");
  	for (a <- simplified_associaton_edges)
  	{
  		println(a);
  	}
  	println("############ SIMPLIFIED DEPENDENCIES: ############");
  	for (d <- simplified_dependency_edges)
  	{
  		println(d);
  	}
  	println("############ CREATING DOT FILE ############");
  
  	// Now write away the results as a DOT graphical description language file.
    str dotGraph = 
    "digraph classes {
         '  fontname = \"Bitstream Vera Sans\"
         '  fontsize = 8
         '  node [ fontname = \"Bitstream Vera Sans\" fontsize = 8 shape = \"record\" ]
         '  edge [ fontname = \"Bitstream Vera Sans\" fontsize = 8 ]
         '
         '  <for (cl <- classes(m)) { /* a for loop in a string template, just like PHP */>
         ' \"N<cl>\" [label=\"{<cl.path[1..] /* a Rascal expression between < > brackets is spliced into the string */>||}\"]
         '  <} /* this is the end of the for loop */>
         '
         '  <for (<from, to> <- generalizations) {>
         '  \"N<to>\" -\> \"N<from>\" [arrowhead=\"empty\"]<}>
         '
         '  <for (<from, to> <- realizations) {>
         '  \"N<to>\" -\> \"N<from>\" [arrowhead=\"empty\", style=\"dotted\"]<}>
         '
         '  <for (<from, to> <- nestedclasses) {>
         '  \"N<to>\" -\> \"N<from>\" [arrowhead=\"odot\"]<}>
         '
         '  <for (<from, to> <- simplified_associaton_edges) {>
         '  \"N<from>\" -\> \"N<to>\" [arrowhead=\"vee\"]<}>
         '
         '  <for (<from, to> <- simplified_dependency_edges) {>
         '  \"N<from>\" -\> \"N<to>\" [arrowhead=\"vee\", style=\"dotted\"]<}>
         '}";
         
         writeFile(|home:///<m.id.authority>.dot|, dotGraph);
         println("DOT file written! Look in your user folder, most likely.");
}